from utils.api_request import *
from utils.data_type import QueryResult, QueryResultItem, wrap_search_results
from utils.llm_client import llm_chat
from agents.paper_search_agent import search_paper_loop
from prompts.detail_outline import DETAIL_OUTLINE_GENERATE_PROMPT
import asyncio
import json
from utils.token_count import glm_tokenize_text
import re
from collections import defaultdict


class DetailOutlineAgent:
    def __init__(self, summary_type: str, origin_query: str, origin_result: dict, paper_search_result: QueryResult = None):
        self.paper_search_result = paper_search_result
        self.summary_type = summary_type
        self.origin_query = origin_query
        self.origin_result = origin_result
        self.detail_outline_str = ""
        self.reference_mark_str = ""
        self.detail_outline_dict = {}
        if self.paper_search_result:
            self.paper_id_list = [item.entity.paper_id for item in self.paper_search_result.results]
            print(f"一共有 {len(self.paper_id_list)} 篇相关文献！")
            # 使用异步方法获取所有论文内容
            self.all_papers_content = asyncio.run(self.get_all_papers_content())

    async def get_entire_paper_content(self, paper_id):
        result = await async_query_by_paper_id(paper_id=paper_id)
        return {paper_id: result}

    async def get_all_papers_content(self):
        # 并发获取所有论文内容
        tasks = [self.get_entire_paper_content(paper_id) for paper_id in self.paper_id_list]
        results = await asyncio.gather(*tasks)
        return results

    def generate_result(self):
        # with open('detail_outline.txt', 'r', encoding='utf-8') as f:
        #     result = f.read()
        paper_chunks = []
        index = 1
        for paper_id, chunk_content in self.origin_result.items():
            temp = ""
            temp += f"<论文块{index}>\n"
            temp += f"<论文ID>{paper_id}</论文ID\n"
            temp += f"<论文标题>{chunk_content['entity']['paper_title']}</论文标题>\n"
            temp += f"<论文块内容>{chunk_content['entity']['chunk_text']}</论文块内容>\n"
            temp += f"</论文块{index}>\n"
            index += 1
            paper_chunks.append(temp)
        paper_chunks_str = "\n".join(paper_chunks)
        prompt = DETAIL_OUTLINE_GENERATE_PROMPT.format(user_requirement=self.origin_query, summary_type=self.summary_type,paper_chunks=paper_chunks_str)
        # with open('prompt.txt','w') as f:
        #     f.write(prompt)
        print(f"token消耗为：{glm_tokenize_text(prompt)}")
        response = llm_chat(messages=[{"role": "user", "content": prompt}],model_name="GLM_4_PLUS",temperature=0)
        result = response.choices[0].message.content
        print(result)

        # 提取 <detail_outline> 和 <reference_mark> 标签内的内容
        detail_outline_pattern = r'<detail_outline>(.*?)</detail_outline>'
        detail_outline_matches = re.findall(detail_outline_pattern, result, re.DOTALL)
        reference_mark_pattern = r'<reference_mark>(.*?)</reference_mark>'
        reference_mark_matches = re.findall(reference_mark_pattern, result, re.DOTALL)

        if detail_outline_matches:
            self.detail_outline_str = detail_outline_matches[0].strip()
        else:
            print("未找到 <detail_outline> 标签内容！")
            self.detail_outline_str = ""

        if reference_mark_matches:
            self.reference_mark_str = reference_mark_matches[0].strip()
        else:
            print("未找到 <reference_mark> 标签内容！")
            self.reference_mark_str = ""

    def parse_detail_outline(self, outline_str):
        """
        解析 detail_outline 字符串为嵌套字典
        支持多级章节
        """
        outline_dict = {}
        current_sections = []
        for line in outline_str.strip().split('\n'):
            line = line.strip()
            if not line:
                continue
            # 确定层级
            header_match = re.match(r'^(#+)\s+(.*)', line)
            item_match = re.match(r'^-\s+(.*)', line)
            if header_match:
                level = len(header_match.group(1))
                title = header_match.group(2).strip()
                if level == 1:
                    # 一级标题（顶级）
                    outline_dict[title] = {}
                    current_sections = [title]
                elif level == 2:
                    # 二级标题
                    section = title
                    outline_dict[current_sections[0]].setdefault(section, {})
                    current_sections = [current_sections[0], section]
                elif level == 3:
                    # 三级标题
                    subsection = title
                    if len(current_sections) >= 2:
                        outline_dict[current_sections[0]][current_sections[1]].setdefault(subsection, {})
                        current_sections = [current_sections[0], current_sections[1], subsection]
                    else:
                        print(f"警告: 三级标题 '{subsection}' 没有父级二级标题。")
                # 可以根据需要增加更多层级
            elif item_match:
                item = item_match.group(1).strip()
                # 添加到当前的最深层级
                if current_sections:
                    target = outline_dict
                    for sec in current_sections:
                        target = target.get(sec, {})
                    if 'items' not in target:
                        target['items'] = []
                    target['items'].append(item)
                else:
                    print(f"警告: 项目 '{item}' 没有对应的章节。")
        return outline_dict

    def parse_reference_mark(self, ref_mark_str):
        """
        解析 reference_mark 字符串为字典
        """
        ref_mark_str = ref_mark_str.strip()
        # 确保 JSON 格式正确
        try:
            ref_mark = json.loads(ref_mark_str)
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
            # 尝试修复常见问题，如单引号替换为双引号
            ref_mark_str = ref_mark_str.replace("'", '"')
            try:
                ref_mark = json.loads(ref_mark_str)
            except json.JSONDecodeError as e2:
                print(f"修复后仍然无法解析 JSON: {e2}")
                ref_mark = {}
        return ref_mark
    def associate_references(self, outline_dict, ref_mark):
        """
        将参考文献 ID 关联到大纲字典中
        递归地遍历大纲字典，并根据章节名称直接添加 references
        """

        def recursive_associate(current_dict):
            for key, value in current_dict.items():
                if key == 'items' or key == 'references':
                    continue  # 跳过 'items' 和 'references' 键
                # 使用当前键作为 path_str 直接匹配 ref_mark
                path_str = key.strip()
                # 查找参考文献
                refs = ref_mark.get(path_str, [])
                # 检查 current_dict[key] 是否为 dict
                if isinstance(current_dict[key], dict):
                    current_dict[key]['references'] = refs
                else:
                    print(f"警告: '{path_str}' 对应的值不是字典，无法添加 'references'")
                    continue
                # 递归处理子章节
                if isinstance(value, dict):
                    recursive_associate(current_dict[key])

        recursive_associate(outline_dict)
        return outline_dict
    def validate_references(self, associated_outline, valid_ids):
        """
        验证关联的参考文献 ID 是否在有效 ID 列表中
        """
        def recursive_validate(current_dict):
            for key, value in current_dict.items():
                if key == 'items':
                    continue
                if 'references' in value:
                    # 过滤无效的 IDs
                    valid_refs = [ref for ref in value['references'] if ref in valid_ids]
                    current_dict[key]['references'] = valid_refs
                if isinstance(value, dict):
                    recursive_validate(value)

        recursive_validate(associated_outline)
        return associated_outline

    def outline_to_markdown(self, outline_dict, level=1):
        """
        将大纲字典转换为 Markdown 格式
        """
        markdown = ""
        for key, value in outline_dict.items():
            if key in ['items', 'references']:
                continue
            markdown += f"{'#' * level} {key}\n\n"
            if 'items' in value:
                for item in value['items']:
                    markdown += f"- {item}\n"
            if 'references' in value and value['references']:
                markdown += "\n**References:**\n"
                for ref in value['references']:
                    markdown += f"- {ref}\n"
            # 递归处理子章节
            sub_markdown = self.outline_to_markdown(value, level + 1)
            if sub_markdown:
                markdown += f"\n{sub_markdown}\n"
        return markdown

    def generate_outline(self):
        self.generate_result()
        # 确保 self.detail_outline_str 和 self.reference_mark_str 已经被赋值
        if not self.detail_outline_str or not self.reference_mark_str:
            print("缺少 detail_outline 或 reference_mark 内容！")
            return

        # 解析 detail_outline
        outline_dict = self.parse_detail_outline(self.detail_outline_str)
        print("解析后的大纲字典：", json.dumps(outline_dict, ensure_ascii=False, indent=4))
        with open('examples_data/outline_dict.json','w') as f:
            f.write(json.dumps(outline_dict, ensure_ascii=False, indent=4))
        # 解析 reference_mark
        ref_mark = self.parse_reference_mark(self.reference_mark_str)
        print("解析后的参考文献标记：", json.dumps(ref_mark, ensure_ascii=False, indent=4))
        with open('examples_data/ref_mark.json','w') as f:
            f.write(json.dumps(ref_mark, ensure_ascii=False, indent=4))
        # 关联参考文献
        associated_outline = self.associate_references(outline_dict, ref_mark)
        self.detail_outline_dict = associated_outline

        print("最终关联后的大纲字典：", json.dumps(self.detail_outline_dict, ensure_ascii=False, indent=4))

        # 验证参考文献 ID 的有效性（可选）
        # 假设您有一个有效的参考文献 ID 列表
        # valid_ids = ["6584feac939a5f4082397b62", "63608e5090e50fcafdee1152", "..."]  # 根据实际情况填充
        # self.detail_outline_dict = self.validate_references(self.detail_outline_dict, valid_ids)

        # 将结果转换为 JSON 并存储
        result_json = json.dumps(self.detail_outline_dict, ensure_ascii=False, indent=4)
        with open('examples_data/parsed_outline.json', 'w', encoding='utf-8') as f:
            f.write(result_json)
        print("解析后的大纲 JSON 已保存到 'parsed_outline.json'。")

        # 将大纲转换为 Markdown 并保存
        markdown_output = self.outline_to_markdown(self.detail_outline_dict)
        with open('parsed_outline.md', 'w', encoding='utf-8') as f:
            f.write(markdown_output)
        print("转换后的 Markdown 输出已保存到 'parsed_outline.md'。")


if __name__ == '__main__':
    origin_query = "What is the current status and challenges of Text2SQL research?"
    user_summary_type = "A review of the current research status in a certain direction"

    # 从文件读取 origin_paper_search_result
    try:
        with open('examples_data/origin_paper_search_result.json', 'r', encoding='utf-8') as f:
            origin_result = json.load(f)
    except FileNotFoundError:
        print("文件 'origin_paper_search_result.json' 未找到，请确保文件存在。")
        origin_result = {}

    # 创建 DetailOutlineAgent 实例
    detail_outline_agent = DetailOutlineAgent(
        summary_type=user_summary_type,
        origin_query=origin_query,
        origin_result=origin_result
    )

    # 生成并解析大纲
    detail_outline_agent.generate_outline()
